/*
 *
 *  * Copyright (c) iwindplus Technologies Co., Ltd.2024-2030, All rights reserved.
 *
 *
 */

package com.iwindplus.base.web.filter;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.core.util.ArrayUtil;
import com.fasterxml.jackson.databind.JsonNode;
import com.iwindplus.base.domain.constant.CommonConstant.ExceptionConstant;
import com.iwindplus.base.domain.constant.CommonConstant.SymbolConstant;
import com.iwindplus.base.util.JacksonUtil;
import com.iwindplus.base.util.UrlUtil;
import com.iwindplus.base.web.domain.property.FilterProperty;
import com.iwindplus.base.web.domain.property.FilterProperty.XssFilterConfig;
import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;

/**
 * 参数去除空格，实现对request body的二次读取，防止夸脚本工具攻击.
 *
 * @author zengdegui
 * @since 2020/4/3
 */
@Slf4j
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private final FilterProperty property;

    public XssHttpServletRequestWrapper(HttpServletRequest request, FilterProperty property) {
        super(request);
        this.property = property;
    }

    @Override
    public String getQueryString() {
        return clean(super.getQueryString());
    }

    @Override
    public String getParameter(String name) {
        return clean(super.getParameter(name));
    }

    @Override
    public String[] getParameterValues(String name) {
        final String[] parameterValues = super.getParameterValues(name);
        if (ArrayUtil.isEmpty(parameterValues)) {
            return new String[0];
        }

        return Arrays.stream(parameterValues)
            .filter(Objects::nonNull)
            .map(this::clean)
            .toArray(String[]::new);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameterMap = super.getParameterMap();
        if (MapUtil.isEmpty(parameterMap)) {
            return Map.of();
        }

        return parameterMap.entrySet().stream()
            .filter(entry -> CharSequenceUtil.isNotBlank(entry.getKey()) && ArrayUtil.isNotEmpty(entry.getValue()))
            .collect(Collectors.toMap(
                Map.Entry::getKey,
                entry -> Arrays.stream(ArrayUtil.isNotEmpty(entry.getValue()) ? entry.getValue() : new String[0])
                    .filter(Objects::nonNull)
                    .map(this::clean)
                    .toArray(String[]::new)
            ));
    }

    @Override
    public Enumeration<String> getParameterNames() {
        return getEnumeration(super.getParameterNames());
    }

    @Override
    public Object getAttribute(String name) {
        return cleanValue(super.getAttribute(name));
    }

    @Override
    public Enumeration<String> getAttributeNames() {
        return getEnumeration(super.getAttributeNames());
    }

    @Override
    public String getHeader(String name) {
        return clean(super.getHeader(name));
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        return getEnumeration(super.getHeaders(name));
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        if (!isJsonRequest()) {
            return super.getInputStream();
        }

        String json = getJson(super.getInputStream());
        if (CharSequenceUtil.isBlank(json)) {
            return super.getInputStream();
        }

        byte[] jsonBytes = parseJsonToBytes(json);
        return getServletInputStream(jsonBytes);
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream(), StandardCharsets.UTF_8));
    }

    private Enumeration<String> getEnumeration(Enumeration<String> headers) {
        List<String> list = new ArrayList<>();
        while (headers.hasMoreElements()) {
            String value = headers.nextElement();
            if (CharSequenceUtil.isNotBlank(value)) {
                list.add(clean(value));
            }
        }
        return Collections.enumeration(list);
    }

    private ServletInputStream getServletInputStream(byte[] jsonBytes) {
        ByteArrayInputStream bis = new ByteArrayInputStream(jsonBytes);
        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return true;
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public int available() {
                return jsonBytes.length;
            }

            @Override
            public void setReadListener(ReadListener readListener) {
                // Do nothing
            }

            @Override
            public int read() {
                return bis.read();
            }
        };
    }

    private boolean isJsonRequest() {
        String header = getHeader(HttpHeaders.CONTENT_TYPE);
        return CharSequenceUtil.isNotBlank(header) && MediaType.APPLICATION_JSON_VALUE.equalsIgnoreCase(header);
    }

    private String getJson(ServletInputStream inputStream) {
        StringBuilder stringBuilder = new StringBuilder();
        try (InputStreamReader inputStreamReader = new InputStreamReader(inputStream, StandardCharsets.UTF_8);
            BufferedReader bufferedReader = new BufferedReader(inputStreamReader)) {
            String line;
            while (null != (line = bufferedReader.readLine())) {
                stringBuilder.append(CharSequenceUtil.trim(line));
            }
        } catch (IOException ex) {
            log.error(ExceptionConstant.IO_EXCEPTION, ex);
        }
        return stringBuilder.toString();
    }

    private byte[] parseJsonToBytes(String json) {
        if (CharSequenceUtil.isBlank(json)) {
            return new byte[0];
        }

        JsonNode jsonNode = JacksonUtil.parseTree(json);
        if (Objects.isNull(jsonNode)) {
            return new byte[0];
        }

        JacksonUtil.cleanJsonNode(jsonNode, this::clean);

        return JacksonUtil.toJsonBytes(jsonNode);
    }

    private Object cleanValue(Object value) {
        if (value instanceof String str) {
            return clean(str);
        } else if (value instanceof Map<?, ?> map) {
            return cleanMap(map);
        } else if (value instanceof List<?> list) {
            return cleanList(list);
        }

        return value;
    }

    private Object cleanMap(Map<?, ?> map) {
        if (MapUtil.isEmpty(map)) {
            return Map.of();
        }

        return map.entrySet().stream()
            .filter(entry -> null != entry.getKey() && null != entry.getValue())
            .collect(Collectors.toMap(
                Map.Entry::getKey,
                entry -> cleanValue(entry.getValue())
            ));
    }

    private Object cleanList(List<?> list) {
        if (CollUtil.isEmpty(list)) {
            return List.of();
        }

        return list.stream()
            .filter(Objects::nonNull)
            .map(this::cleanValue)
            .collect(Collectors.toList());
    }

    private String clean(String value) {
        if (value == null) {
            return null;
        }

        if (CharSequenceUtil.isBlank(value)) {
            return SymbolConstant.EMPTY_STR;
        }

        final XssFilterConfig cfg = this.property.getXss();
        // 是否启用跳过
        if (Boolean.TRUE.equals(cfg.getEnabledSkip())) {
            // 忽略的 API
            final List<String> ignoredApi = cfg.getIgnoredApi();
            if (CollUtil.isNotEmpty(ignoredApi) && UrlUtil.isMatchLike(ignoredApi, getRequestURI())) {
                return value;
            }

            // 忽略特定符号
            final List<String> symbols = cfg.getIgnoredSymbol();
            if (CollUtil.isNotEmpty(symbols) && symbols.stream().anyMatch(value::contains)) {
                return value;
            }
        }

        // 真正 XSS 过滤
        return JacksonUtil.cleanByJsoup(value, cfg.getTagWhiteList());
    }
}